# -*- coding: utf-8 -*-
"""
Created by edc on 2020/7/27
"""

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import MaxPooling2D
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dense
from tensorflow.keras import backend as K

@staticmethod
def build(width, height, depth, classes):
    """

    @param width: 图片宽
    @param height: 图片高
    @param depth: 图片深
    @param classes: 分类数目
    @return: 未build的LeNet模型
    """
    model = Sequential()
    # 确保深度是在后面的
    if K.image_data_format() == "channels_first":
        inputShape = (depth, height, width)

    inputShape = (height, width, depth)

    model.add(Conv2D(20, (5, 5), padding="same", input_shape=inputShape))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(Conv2D(50, (5, 5), padding="same"))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    model.add(Flatten())
    model.add(Dense(500))
    model.add(Activation("relu"))

    model.add(Dense(classes))
    model.add(Activation("softmax"))

    return model
